package com.aloudmeta.grpc.support.validation;

import io.envoyproxy.pgv.ValidationException;
import io.envoyproxy.pgv.ValidatorIndex;
import io.grpc.*;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.protocol.grpc.interceptors.ServerInterceptor;

import static org.apache.dubbo.common.constants.CommonConstants.PROVIDER;

/**
 * huaixin 2022/2/11 10:49 AM
 */
@Activate(group = {PROVIDER}, order=5300)
public class ValidationServerInterceptor implements ServerInterceptor {
    private ValidatorIndex validatorIndex;

    public void setValidatorIndex(ValidatorIndex validatorIndex) {
        this.validatorIndex = validatorIndex;
    }

    @Override
    public <ReqT, RespT> ServerCall.Listener<ReqT> interceptCall(ServerCall<ReqT, RespT> call, Metadata headers, ServerCallHandler<ReqT, RespT> next) {
        return new ForwardingServerCallListener.SimpleForwardingServerCallListener<ReqT>(next.startCall(call, headers)) {

            // Implementations are free to block for extended periods of time. Implementations are not
            // required to be thread-safe.
            private boolean aborted = false;

            @Override
            public void onMessage(ReqT message) {
                try {
                    validatorIndex.validatorFor(message.getClass()).assertValid(message);
                    super.onMessage(message);
                } catch (ValidationException ex) {
                    StatusRuntimeException status = ValidationExceptions.asStatusRuntimeException(ex);
                    aborted = true;
                    call.close(status.getStatus(), status.getTrailers());
                }
            }

            @Override
            public void onHalfClose() {
                if (!aborted) {
                    super.onHalfClose();
                }
            }
        };
    }
}
